import json
import os
import traceback

from mitmproxy import http

#import flowUtil

# 自定义上下文标识符
TEMU_SELF_HEADER_FLAG = "anti-content"

CONFIG_FILE_PATH = "config.json"
# 配置常量
CA_CERT_NAME = "mitmproxy-ca-cert.cer"

# 店小秘
DXM_AUTH_URL = "https://www.dianxiaomi.com/api/userInfo.json"

# TEMU 处理
TEMU_HOST = "https://seller.kuajingmaihuo.com"
TEMU_QUAN_TUO_AUTH_URL = "https://seller.kuajingmaihuo.com/bg/quiet/api/mms/userInfo"

# TEMU 半托认证url
TEMU_BANTUO_AUTH_URL = "https://www.dianxiaomi.com/shop/list/pddkj.htm"
# TEMU_BANTUO_AUTH_URL = "https://agentseller.temu.com/api/seller/auth/userInfo"

# shein 半托认证url
SHEIN_BANTUO_AUTH_URL = "https://www.dianxiaomi.com/sheinProduct/add.htm"

class JsonModifier:
    """HTTP响应修改工具"""

    def request(self, flow: http.HTTPFlow):
        # print(f"收到请求：{flow.request.url}")
        # if flow.server_conn.timestamp_start is None:
        #     print(f"请求开始时间未设置，正在设置...{flow.request.url}")
        #     flow.server_conn.via = ServerSpec(("https",f"{proxy_ip}:{proxy_port}"))
        #
        #     # 如果上游代理需要认证
        #     if proxy_username and proxy_password:
        #         credentials = f"{proxy_username}:{proxy_password}"
        #         encoded_credentials = base64.b64encode(credentials.encode()).decode()
        #         flow.request.headers["Proxy-Authorization"] = f"Basic {encoded_credentials}"

        is_temu_request = False
        # 设置TEMU本站上下文标志
        if TEMU_SELF_HEADER_FLAG in flow.request.headers:
            flow.__dict__["is_temu_request"] = 1
            is_temu_request = True
            print(f"识别到TEMU本站请求，跳过处理")
            return flow

        if is_temu_request:
            pass
        else:
            """还原请求头"""
            if flow.request.url.startswith(TEMU_HOST) or flow.request.url.startswith("https://agentseller.temu.com/"):
                print("检测到大卖家TEMU 请求，正在处理...")
                if "mallid" in flow.request.headers and TEMU_SELF_HEADER_FLAG not in flow.request.headers:
                    header_mallid = flow.request.headers["mallid"]

                    for original_id_str, new_id in self.mapping.items():
                        try:
                            if new_id == header_mallid:
                                flow.request.headers["mallid"] = original_id_str
                                print(f"mallid 还原成功: {header_mallid} -> {original_id_str}")
                                break
                        except Exception as e:
                            traceback.print_exc()
                            print(f"大卖家TEMU请求处理失败: {str(e)}")
                    print(f"还原后mallid：{flow.request.headers['mallid']}")

    def __init__(self):
        """初始化时加载配置文件"""
        self.config = self._load_config()
        self.mapping = self.config.get('mallid_mapping', {})

    def _load_config(self):
        """加载外部配置文件"""
        try:
            if os.path.exists(CONFIG_FILE_PATH):
                with open(CONFIG_FILE_PATH, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    print("✅ 配置文件加载成功")
                    return config
            else:
                raise FileNotFoundError(f"配置文件不存在: {CONFIG_FILE_PATH}")
        except Exception as e:
            print(f"❌ 配置文件加载失败: {str(e)}")
            raise e

    def response(self, flow: http.HTTPFlow):
        response_url = flow.request.url
        # 检查上下文标志
        if getattr(flow, "is_temu_request", 0) == 1:
            print("💡 检测到TEMU本站上下文标志，跳过处理")
            return
        else:
            # print(f"处理响应：{response_url}")
            # if response_url.__contains__("holly/mall/manage/list"):
            #     print(f"命中测试url:{response_url}")

            self.add_temu_quantuo_rule(flow)
            self.add_temu_bantuo_rule(flow)
            self.add_shein_bantuo_rule(flow)
            self.add_temu_ads_rule(flow)

            self.add_meiguo_rule(flow)

            # self.add_dxm_rule(flow)

    def add_temu_quantuo_rule(self, flow):
        # 全托
        if flow.request.url == TEMU_QUAN_TUO_AUTH_URL:
            print(f"🚀 命中TEMU全托拦截:{flow.request.url}")
            try:
                # 解析并修改JSON响应
                response = json.loads(flow.response.text)
                response_text = flow.response.text

                # 安全访问嵌套字段
                if (response.get("result", {}).get("companyList") and
                        response["result"]["companyList"][0].get("malInfoList") and
                        len(response["result"]["companyList"][0]["malInfoList"]) > 0):
                    pass
                    # # 获取原始mallId并转换为字符串进行匹配
                    # original_mallid = response["result"]["companyList"][0]["malInfoList"][0]["mallId"]
                    #
                    # # 只有在映射表中存在时才替换，否则保留原始值
                    # if str(original_mallid) in self.mapping:
                    #     new_mallid = self.mapping[str(original_mallid)]
                    #     response["result"]["companyList"][0]["malInfoList"][0]["mallId"] = new_mallid
                    #     print(f"✅ 映射成功: {original_mallid} → {new_mallid}")
                    # else:
                    #     print(f"ℹ️ 未找到匹配，保留原始值: {original_mallid}")

                else:
                    # 没有数据，则组装新数据
                    mock_json_data = '{"success":true,"errorCode":1000000,"errorMsg":null,"result":{"userId":10427381982257,"maskMobile":"+86 153****7996","maskOwnerMobile":"7996","accountType":1,"accountStatus":1,"tokenType":1,"companyList":[{"companyName":"-","malInfoList":[{"mallId":634418213187299,"mallName":"ouzhiluckyshop","logo":"https://img.cdnfe.com/supplier-public-tag/1f29b0ce60/466383dd-e0f6-49e2-b93c-43fcc780d78b_300x300.jpeg","mallStatus":1,"isSemiManagedMall":false},{"mallId":634418213741919,"mallName":"delilucluckshop","logo":"https://img.cdnfe.com/supplier-public-tag/1f29b0a1b0/93042c96-726f-4798-8b1b-8623ca95ce1a_300x300.jpeg","mallStatus":1,"isSemiManagedMall":false},{"mallId":634418213742069,"mallName":"DreamHomes","logo":"https://img.cdnfe.com/supplier-public-tag/1f135083f38/6387cc55-d7e8-4a58-8edc-1510c63219e1_300x300.jpeg","mallStatus":1,"isSemiManagedMall":true},{"mallId":634418216135768,"mallName":"OutdoorNestFurniture","logo":"https://img.cdnfe.com/supplier-public-tag/1f23584280/0b5bfff1-796b-4825-aa28-5b3e0d37c9f8_1000x1000.jpeg","mallStatus":1,"isSemiManagedMall":true},{"mallId":634418216135828,"mallName":"GreenGardenFurniture","logo":"https://img.cdnfe.com/supplier-public-tag/1f23584280/f6f376e8-b9c9-4075-9058-9369471688de_1000x1000.jpeg","mallStatus":1,"isSemiManagedMall":true}]}],"nickName":null,"roleNameList":["主账号"]}}'
                    print("mock数据")
                    response_text = mock_json_data

                # 遍历所有映射项，尝试匹配并替换
                for original_id_str, new_id in self.mapping.items():
                    try:
                        if original_id_str not in response_text:
                            continue
                        response_text = response_text.replace(original_id_str, new_id)
                        print(f"替换TEMU全托管成功: {original_id_str} -> {new_id}")
                    except ValueError:
                        print(f"⚠️ 非法ID格式: {original_id_str}")

                    # 移除压缩头防止解码错误
                    if "Content-Encoding" in flow.response.headers:
                        del flow.response.headers["Content-Encoding"]

                flow.response.text = response_text



            except (IndexError, KeyError, json.JSONDecodeError) as e:
                print(f"⚠️ TEMU全托JSON修改错误: {str(e)}")

    def add_shein_bantuo_rule(self, flow: http.HTTPFlow):
        if flow.request.url == "https://www.dianxiaomi.com/userTemplate/sheinAdd.htm?shopId=":
            print(f"🚀 命中SHEIN半托拦截:{flow.request.url}")

            response_text = flow.response.text
            for original_id_str, new_id in self.mapping.items():
                try:
                    if original_id_str not in response_text:
                        continue
                    response_text = response_text.replace(original_id_str, new_id)
                    print(f"替换shein半托管成功: {original_id_str} -> {new_id}")
                except ValueError:
                    print(f"⚠️ 非法ID格式: {original_id_str}")
            flow.response.text = response_text

    def add_temu_bantuo_rule(self, flow: http.HTTPFlow):
        if flow.request.url in self.config['TEMU']['auth_url']:
            print(f"🚀 命中TEMU半托拦截:{flow.request.url}")

            response_text = flow.response.text
            for original_id_str, new_id in self.mapping.items():
                try:
                    if original_id_str not in response_text:
                        continue
                    response_text = response_text.replace(original_id_str, new_id)
                    print(f"替换TEMu半托管成功: {original_id_str} -> {new_id}")
                except ValueError:
                    print(f"⚠️ 非法ID格式: {original_id_str}")

            flow.response.text = response_text

    def add_dxm_rule(self, flow):
        if flow.request.url == DXM_AUTH_URL:
            print(f"🚀 命中店小秘拦截:{flow.request.url}")

            response_text = flow.response.text
            for original_id_str, new_id in self.mapping.items():
                try:
                    if original_id_str not in response_text:
                        continue
                    response_text = response_text.replace(original_id_str, new_id)
                    print(f"替换店小秘成功: {original_id_str} -> {new_id}")
                except ValueError:
                    print(f"⚠️ 非法ID格式: {original_id_str}")

            flow.response.text = response_text

    def test(self):
        for original_id_str, new_id in self.mapping.items():
            response_text = "aaa123"
            if original_id_str not in response_text:
                continue
            try:
                response_text = response_text.replace(original_id_str, new_id)
                print(f"替换TEMU全托管成功: {original_id_str} -> {new_id}")
            except ValueError:
                print(f"⚠️ 非法ID格式: {original_id_str}")

            print(f"替换结果:{response_text}")

    # 批量添加广告功能
    def add_temu_ads_rule(self, flow):
        if flow.request.host.endswith("ads.temu.com") and flow.request.path=='/api/v1/coconut/mall/query_mall_detail':
        # if flow.request.url == "https://us.ads.temu.com/api/v1/coconut/mall/query_mall_detail":
            print(f"🚀 命中temu批量添加广告功能拦截:{flow.request.url}")

            # 固定用这个店铺id
            FIX_MALL_ID = "634418215979225"

            origin_text = flow.response.text
            try:
                if "success" in origin_text:
                    json_data = flow.response.json()
                    if json_data['result'] and 'mall_id' in json_data['result']:
                        json_data['result']['mall_id'] = FIX_MALL_ID

                    flow.response.text = json.dumps(json_data)
                    #flowUtil.print_log(flow)
                else:
                    print(f"temu 批量添加广告功能认证接口返回失败:{flow.response.text}")
            except Exception as e:
                print("X temu批量添加广告功能处理失败")
                traceback.print_exc()

    def add_meiguo_rule(self, flow):
        if flow.request.url == "https://seller.temu.com/home.html":
            print(f"🚀 命中美国功能拦截:{flow.request.url}")

            response_text = flow.response.text
            for original_id_str, new_id in self.mapping.items():
                try:
                    if original_id_str not in response_text:
                        continue
                    response_text = response_text.replace(original_id_str, new_id)
                    print(f"替换美国功能成功: {original_id_str} -> {new_id}")
                except ValueError:
                    print(f"⚠️ 非法ID格式: {original_id_str}")

            flow.response.text = response_text

# 初始化插件
addons = [
    JsonModifier()
]